from path import *
from module import *
from mpc import *

from matplotlib import pyplot as plt

wayPoints = np.array([[0, 0], [1, -3], [2, 2], [5, 0]])
path = Path(wayPoints)
path.interplot(100)
iPath = path.getInterplotPoints()

# plt.plot(iPath[:, 2], iPath[:, 3])

car = Module(5, -5, 1, -1, 0.5)

mpc = MPC(10, 4, 2)
mpc.setQ(np.diag([0, 0, 100, 100]))
mpc.setR(np.eye(2) * 1)
mpc.setStateRange(car.getStateRange()[0], car.getStateRange()[1])
mpc.setInputRange(car.getControlRange()[0], car.getControlRange()[1])

timestep = 0.1

states: list[list[float]] = []
for i in range(500):
    state = car.getState()
    theta = state[0]
    v = state[1]
    x = state[2]
    y = state[3]
    states.append([float(theta), float(v), float(x), float(y)])
    curPoint = path.getNearestPointIndex(x, y)
    print(f"i:{curPoint},v({x},{y}),p({iPath[curPoint][2]},{iPath[curPoint][3]})")
    A = []
    B = []
    ref: list[float] = []
    for i in range(curPoint, curPoint + 11):
        ref.append(iPath[i][0])
        ref.append(iPath[i][1])
        ref.append(iPath[i][2])
        ref.append(iPath[i][3])
    for i in range(curPoint, curPoint + 10):
        A.append(car.getMatrixA(iPath[i][1], iPath[i][0], timestep))
        B.append(car.getMatrixB(timestep))
    xRef = np.array(ref).reshape(-1, 1)
    mpc.setA(A)
    mpc.setB(B)
    x, u = mpc.solve(xRef, car.getState())
    w, a = u[0, 0], u[0, 1]
    car.update(w, a, timestep)
ax = plt.subplot(1, 2, 1)
ax.plot(iPath[:, 2], iPath[:, 3], label="path")
ax.scatter(iPath[:, 2], iPath[:, 3])
theta = [row[0] for row in states]
v = [row[1] for row in states]
x = [row[2] for row in states]
y = [row[3] for row in states]
ax.plot(x, y, label="real")
ax.scatter(x, y)
ax.legend()
bx = plt.subplot(1, 2, 2)
bx.plot(theta, label="theta")
bx.plot(v, label="v")
# bx.plot(x, label="x")
# bx.plot(y, label="y")
bx.legend()
plt.show()
